import numpy as np

from scipy.optimize import linear_sum_assignment
import json

def calculate_cost_matrix(boxes1, boxes2):
    """
    计算两组3D框之间的成本矩阵，这里使用欧几里得距离作为成本
    :param boxes1: 第一组3D框，形状为 (m, 3)
    :param boxes2: 第二组3D框，形状为 (n, 3)
    :return: 成本矩阵，形状为 (m, n)
    """
    m = boxes1.shape[0]
    n = boxes2.shape[0]
    cost_matrix = np.zeros((m, n))
    for i in range(m):
        for j in range(n):
            # 计算两个3D框中心的欧几里得距离
            cost_matrix[i, j] = np.linalg.norm(boxes1[i][:3] - boxes2[j][:3])
    return cost_matrix


def match_3d_boxes(boxes1, boxes2):
    """
    使用匈牙利算法对两组3D框进行匹配
    :param boxes1: 第一组3D框，形状为 (m, 3)
    :param boxes2: 第二组3D框，形状为 (n, 3)
    :return: 匹配结果，形状为 (k, 2)，其中 k 为匹配的对数
    """
    boxes1=np.array(boxes1)
    boxes2=np.array(boxes2)
    # 计算成本矩阵
    cost_matrix = calculate_cost_matrix(boxes1, boxes2)
    # 使用匈牙利算法求解最小匹配
    row_indices, col_indices = linear_sum_assignment(cost_matrix)
    matches = np.column_stack((row_indices, col_indices))
    return matches

def parse_box_string_to_array(box_str):
    # 同 parse_box_string，但返回 shape=(N,9) 的 numpy 数组
    s = box_str.replace("tensor(", "").replace(")", "")
    s = s.replace("[", "").replace("]", "").replace(",", "")
    nums = list(map(float, re.sub(r'\s+', ' ', s).strip().split()))
    arr = np.array(nums).reshape(-1, 9)
    return arr

def parse_gt_names(gt_str):
    # 去掉首尾的方括号和多余空格
    s = gt_str.strip()[1:-1].strip()
    # 用正则匹配所有带引号的子串
    return re.findall(r"'(.*?)'", s)   


gt=json.load(open('gt.json',"rb"))

pre=json.load(open('pre.json',"rb"))

# 类别映射关系
category_mapping = {
    "car": 0,
    "truck": 1,
    "forklift": 2,
    "tractor": 3,
    "flatbed": 4,
    "trailer": 5,
    "bicycle": 6,
    "cyclist": 7,
    "pedestrian": 8,
    "traffic_cone": 9,
    "water_horse": 10,
    "guardrail": 11,
    "unhnown": 12,
}

TP=np.zeros(7)
FP=np.zeros(7)
FN=np.zeros(7)

label_err=0
dis_err=0
        
for i in range(len(gt)):
    gt_boxes=gt[i]['box3d_lidar']
    gt_labels=gt[i]['label_preds']
    pre_boxes=pre[i]['box3d_lidar']
    pre_labels=pre[i]['label_preds']

    matches=match_3d_boxes(gt_boxes,pre_boxes)

    fp_obj=np.arange(len(pre_boxes))
    fp_obj=np.setdiff1d(fp_obj, matches[:,0])

    for i in fp_obj:
        FP[pre_labels[i]]+=1

    fn_obj=np.arange(len(gt_boxes))
    fn_obj=np.setdiff1d(fn_obj, matches[:,1])

    for i in fn_obj:
        FN[int(gt_labels[i])]+=1

    
    for match in matches:
        gt_box=np.array(gt_boxes[match[0]])
        pre_box=np.array(pre_boxes[match[1]])

        gt_label=gt_labels[match[0]]
        pre_label=pre_labels[match[1]]

        tp_flag=True

        # if gt_label!=pre_label:
        #     tp_flag=False
        #     FP[pre_label]+=1
        #     FN[int(gt_label)]+=1
        #     label_err=label_err+1
        #     continue
        center_dis=np.linalg.norm(gt_box[:3] - pre_box[:3])
        if center_dis>4:
            tp_flag=False
            FP[pre_label]+=1
            FN[pre_label]+=1
            dis_err=dis_err+1
            continue

        TP[pre_label]+=1


print('clsname   TP   FP   FN   precision  recall')
for i in range(len(class_names)):
    print('%s    %d   %d   %d    %.2f       %.2f' % (class_names[i],TP[i],FP[i],FN[i],TP[i]/(TP[i]+FP[i]+1e-6),TP[i]/(TP[i]+FN[i]+1e-6)))
        